import inspect
from typing import Any

from langchain.prompts import StringPromptTemplate
from pydantic import BaseModel
from pydantic.class_validators import validator


def get_source_code(funtion_name):
    return inspect.getsource(funtion_name)


class FunctionExplainerPromptTemplate(StringPromptTemplate, BaseModel):
    """
    A custom prompt template that takes in the function name as input,
     and formats the prompt template to provide the source code of the function
    """
    @validator("input_variables")
    def validate_input_variables(cls, v):
        """
        Validate that the input variables are correct.
        :param v:
        :return:
        """
        if len(v) != 1 or "function_name" not in v:
            raise ValueError("funtion_name must be be the only input_variable.")
        return v

    def format(self, **kwargs: Any) -> str:
        """
        Get the source code of the function
        :param kwargs:
        :return:
        """
        # Get the source code of the function
        source_code = get_source_code(kwargs["function_name"])
        # Generate the prompt to be sent to the language model
        prompt = f"""
        Given the function name and source code, generate an English language explanation of the function.
        Function Name: {kwargs["function_name"].__name__}
        Source Code:
        {source_code}
        Explanation:
        """
        return prompt

    def _prompt_type(self) -> str:
        return "function-explainer"


if __name__ == '__main__':
    fn_explainer = FunctionExplainerPromptTemplate(input_variables=["function_name"])

    # 为函数"get_source_code"生成一个提示
    prompt = fn_explainer.format(function_name=get_source_code)
    print(prompt)
